/*! 
  @file   MarkerDetector.cpp
  @brief  Find markers from captured image
  @author Kyohei Otsu

  Find red border from an image and publish inner square
*/

#include <cassert>
#include <vector>

#include <ros/ros.h>
#include <nodelet/nodelet.h>
#include <cv_bridge/cv_bridge.h>
#include <image_transport/camera_subscriber.h>
#include <image_transport/image_transport.h>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <sensor_msgs/image_encodings.h>

#include "camera/MarkerImage.h"


using namespace std;
using namespace cv;
using namespace camera;

#define FOR(x, a, b) for(typeof(a) x = (a); x < static_cast<typeof(a)>(b); x++)
#define FORIT(x, v) for(typeof(v.begin()) x = v.begin(); x != v.end(); x++)

typedef std::vector< std::vector<cv::Point> > Contours;

namespace marker_detector
{

const int MARKER_SIZE = 50;


//! @brief Detector of Marker Contour
class MarkerDetector : public nodelet::Nodelet
{

private:

  image_transport::Subscriber img0_sub_; //!< image subscriber
  image_transport::Subscriber img1_sub_; //!< image subscriber
  ros::Publisher marker_pub_;           //!< marker image publiser

public:

  //! called when initialization
  virtual void onInit();

protected:

  //! @brief  Extract red color
  //! @param  src source image (CV_8UC3)
  //! @param  dst destination image (CV_8UC1)
  //! @param  color should be always 0 (=red)
  void binalize_by_color(cv::Mat &src, cv::Mat &dst, int color=0);

  //! @brief  detect marker border and publish inner image
  int detect_and_publish(const sensor_msgs::ImageConstPtr& image_msg);

  //! @brief  callback function
  void imageCallback(const sensor_msgs::ImageConstPtr& image_msg);

 private:

  //! @brief  Approximate contour
  //! @param  src vector of corner points
  //! @param  brief  dst vector of corner points after approximation
  int _approx(std::vector<cv::Point> &src, std::vector<cv::Point2f> &dst);

  //! @brief  Calc perimeter of contour
  //! @param  curve vector of corner points 
  //! @return length of perimether
  float _calc_perimeter(std::vector<cv::Point2f> &curve);

  //! @brief  Calc position of the marker
  //! @param  c vector of corner points 
  //! @param  w image width
  int _calc_angle(std::vector<cv::Point2f> &c, int w);

  //! @brief  sort corners in a marker
  //! @param  c vector of corner points 
  int _sort_corners(std::vector<cv::Point2f>&c);
  std::vector<cv::Point2f>::iterator _sort_corners_find(
      std::vector<cv::Point2f> &c, int n, cv::Point2f &gc);
};


void MarkerDetector::onInit()
{
  NODELET_DEBUG("Initializing marker detector...");
  ros::NodeHandle &nh = getNodeHandle();
  //ros::NodeHandle &private_nh = getPrivateNodeHandle();
  image_transport::ImageTransport it(nh);

  // Read params

  // subscriber
  img0_sub_ = it.subscribe("imageL", 5, &MarkerDetector::imageCallback, this);
  img1_sub_ = it.subscribe("imageR", 5, &MarkerDetector::imageCallback, this);

  // publisher
  marker_pub_      = nh.advertise<camera::MarkerImage>("marker_img", 1); 
}

void MarkerDetector::imageCallback(const sensor_msgs::ImageConstPtr& image_msg)
{
  detect_and_publish(image_msg);
}

int MarkerDetector::detect_and_publish(const sensor_msgs::ImageConstPtr &image_msg)
{
  cv_bridge::CvImageConstPtr msg;
  msg = cv_bridge::toCvShare(image_msg, sensor_msgs::image_encodings::BGR8);
  Mat src = msg->image;

  cv::imwrite("image.png", src); 
  assert(src.type() == CV_8UC3);

  // preparation for publish
  std::string enc = sensor_msgs::image_encodings::BGR8;
  Mat marker(Size(MARKER_SIZE, MARKER_SIZE), CV_8UC3);
  cv_bridge::CvImagePtr cv_ptr(new cv_bridge::CvImage(msg->header, enc, marker));

  MarkerImagePtr marker_img(new MarkerImage);
  marker_img->header = msg->header;
  marker_img->pos.camera_id = (msg->header.frame_id == "/camera0")? 0: 1;

  // preparker_pose ideal square
  vector<Point2f> ideal;
  ideal.push_back(Point2f(marker.cols, 0.0));
  ideal.push_back(Point2f(0.0, 0.0));
  ideal.push_back(Point2f(0.0, marker.rows));
  ideal.push_back(Point2f(marker.cols, marker.rows));

  // number of detected markers
  int n_detected = 0;


  // binarize image by color (=red)
  Mat in_mono(src.size(), CV_8UC1);
  GaussianBlur( src, src, Size( 7, 7 ), 0, 0 ); //blur
  binalize_by_color(src, in_mono, 0);

  // find contours
  Contours contours;
  vector< cv::Vec4i > hi; // hierarchy
  cv::findContours(in_mono, contours, hi, CV_RETR_CCOMP, CV_CHAIN_APPROX_SIMPLE);
  //ROS_DEBUG("num c = %d", contours.size());

  FOR(pid, 0, contours.size()) 
  {
    // check if the contour has child
    int cid = hi[pid][2];
    if(cid < 0) continue;

    // approximation
    vector<Point2f> pcurve;
    vector<Point2f> ccurve;
    _approx(contours[pid], pcurve);
    _approx(contours[cid], ccurve);

    // check if the contour is convex
    ROS_DEBUG("%d --> %d", (int)pcurve.size(), (int)ccurve.size());
    if(pcurve.size() != 4 || ccurve.size() != 4) continue;

    // sort detected marker image points
    _sort_corners(ccurve);

    // check distortion
    float ratio = (ccurve[3].y-ccurve[0].y) / (ccurve[2].y-ccurve[1].y);
    if(ratio > 3 || ratio < 1./3)
    {
      ROS_ERROR("Marker detected, but distortion is large %.2f", ratio);
      continue;
    }

    // check size
    float perimeter = _calc_perimeter(ccurve);
    ROS_ERROR("perimeter: %.2f", perimeter);
    if(perimeter < 1) //TODO no sense
    {
      ROS_ERROR("Marker detected, but it is too small %.2f", perimeter);
      continue;
    }

    // transformation
    // calculate perspective transform
    Mat map = getPerspectiveTransform(ccurve, ideal);
    warpPerspective(src, cv_ptr->image, map, marker.size());

    cv::imwrite("marker.png", cv_ptr->image);

    // publish image and meta data
    cv_ptr->toImageMsg(marker_img->marker);
    marker_img->pos.angle = _calc_angle(ccurve, src.cols);
    ROS_INFO("Marker detected. angle:%d", marker_img->pos.angle);
    marker_pub_.publish(marker_img);

    n_detected++;
  }

  return n_detected;
}

void MarkerDetector::binalize_by_color(cv::Mat &src, cv::Mat &dst, int color)
{
  for(int y = 0; y < src.rows; y++) {
    cv::Vec3b *ptr = src.ptr<cv::Vec3b>(y);
    uchar *ptr_mono = dst.ptr<uchar>(y);
    for(int x = 0; x < src.cols; x++) {
      cv::Vec3b bgr = ptr[x];
      switch(color) {
        case 0: // RED
          if(bgr[2] > 100 && bgr[2] > 1.5*bgr[0] && bgr[2] > 1.5*bgr[1]) ptr_mono[x] = 255;
          else ptr_mono[x] = 0;
          break;
      }
    }
  }
}

int MarkerDetector::_approx(std::vector<cv::Point> &src, std::vector<cv::Point2f> &dst)
{
  Mat _curve(Size(1, src.size()), CV_32FC2);
  FOR(i, 0, src.size()) {
    _curve.at<Point2f>(i, 0) = Point2f(src[i].x, src[i].y);
  }
  approxPolyDP(_curve, dst, 7, true);
  return dst.size();
}

float MarkerDetector::_calc_perimeter(std::vector<cv::Point2f> &curve)
{
  float perimeter = 0.0;
  FORIT(it, curve) {
    typeof(it) next = (it+1!=curve.end())? it+1: curve.begin();
    perimeter += sqrt(pow(it->x - next->x, 2) + pow(it->y - next->y, 2));
  }
  ROS_DEBUG("Perimeter: %f", perimeter);
  return perimeter;
}

int MarkerDetector::_calc_angle(std::vector<cv::Point2f> &c, int w)
{
  float gx = 0;
  const int fov = 60;
  FORIT(i, c) gx += i->x / 4;
  return - (gx / w - 0.5) * fov;
}

std::vector<cv::Point2f>::iterator MarkerDetector::_sort_corners_find(
    std::vector<cv::Point2f> &c, int n, cv::Point2f &gc) 
{
  vector<Point2f>::iterator it;
  for(it = c.begin(); it != c.end(); it++) {
    switch(n) {
      case 0:
        if(it->x >= gc.x && it->y <= gc.y) return it;
        break;
      case 1:
        if(it->x <= gc.x && it->y <= gc.y) return it;
        break;
      case 2:
        if(it->x <= gc.x && it->y >= gc.y) return it;
        break;
      case 3:
        if(it->x >= gc.x && it->y >= gc.y) return it;
        break;
    }
  }
  return it; // c.end()
}

int MarkerDetector::_sort_corners(std::vector<cv::Point2f>&c)
{
  // 1---0
  // | g |
  // 2---3
  Point2f gc(0, 0);
  vector<Point2f>::iterator it;
  for(it = c.begin(); it != c.end(); it++) {
    gc.x += it->x / 4;
    gc.y += it->y / 4;
  }
  if(c[0].x >= gc.x && c[0].y <= gc.y
      && c[1].x <= gc.x && c[1].y <= gc.y
      && c[2].x <= gc.x && c[2].y >= gc.y
      && c[3].x >= gc.x && c[3].y >= gc.y ) {
    return 0;
  } else {
    // sort
    vector<Point2f> tmp = c;
    c.clear();
    for(int i = 0; i < 4; i++) {
      vector<Point2f>::iterator next = _sort_corners_find(tmp, i, gc);
      if(next != tmp.end()) c.push_back(*next);
      else {
        ROS_DEBUG("Corner sort failed!");
        return -1;
      }
    }
    return 1;
  }
}

} // end of namespace


#if 0
int main(int argc, char **argv)
{
  ros::init(argc, argv, "marker_detector");
  marker_detector::MarkerDetector md;
  
  ros::spin();
  return 0;
}
#endif


#include <pluginlib/class_list_macros.h>
PLUGINLIB_DECLARE_CLASS(camera, MarkerDetector, marker_detector::MarkerDetector, nodelet::Nodelet)
